import numpy as np
import SimpleITK as sitk
import matplotlib
import os

matplotlib.use("TkAgg")
import matplotlib.pyplot as plt

image_folder = "data/image"
label_folder = "data/label"

# 目标尺寸
target_size = (128, 256, 256)  # (Z, Y, X)
target_spacing = (1.0, 1.0, 1.0)

# 载入 mha image label
def load_image(filename):
    # 读取图像文件
    image_path = os.path.join(image_folder, filename)
    label_path = os.path.join(label_folder, filename)

    if os.path.exists(image_path) and os.path.exists(label_path):
        image = sitk.ReadImage(image_path)
        label = sitk.ReadImage(label_path)

        # 进行重采样
        resampled_image = resample_image(image, is_label=False)
        resampled_label = resample_image(label, is_label=True)

        # 转换为 NumPy 数组
        img_array = sitk.GetArrayFromImage(resampled_image)
        label_array = sitk.GetArrayFromImage(resampled_label)
        return img_array, label_array, img_array.shape


# 重采样
def resample_image(origin_img, is_label=False, target_size = target_size, target_spacing = target_spacing):

    """对 MHA 图像进行 3D 重采样"""
    original_size = np.array(origin_img.GetSize(), dtype=np.int32)  # 原始尺寸 (X, Y, Z)
    original_spacing = np.array(origin_img.GetSpacing(), dtype=np.float32)  # 原始间距

    # 计算新的像素间距
    new_spacing = original_spacing * (original_size / np.array(target_size[::-1], dtype=np.float32))

    # 创建重采样过滤器
    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(origin_img)
    # 设置目标图像的信息
    resampler.SetOutputSpacing(new_spacing)  # 设定目标间距
    resampler.SetSize(target_size[::-1])  # 设定目标尺寸
    resampler.SetOutputDirection(origin_img.GetDirection())  # 保持方向不变
    resampler.SetOutputOrigin(origin_img.GetOrigin())  # 保持原点位置不变

    # 选择插值方式（最近邻插值用于标签，线性插值用于图像）
    resampler.SetInterpolator(sitk.sitkNearestNeighbor if is_label else sitk.sitkLinear)

    return resampler.Execute(origin_img)


def extract_slice(img_array, label_array, z_index):
    if z_index < 0 or z_index >= img_array.shape[0] or z_index >= label_array.shape[0]:
        raise ValueError(f"z_index {z_index} is out of range. Valid range: 0 to {img_array.shape[0] - 1}")
    return img_array[z_index, :, :], label_array[z_index, :, :]


def plot_slice(img_array, label_array, z_index, img_cmap="grey", label_camp="hot"):
    plt.figure(figsize=(6, 6))
    plt.imshow(img_array, cmap=img_cmap)  # 第一张图像，设置透明度
    plt.imshow(np.ma.masked_where(label_array == 0, label_array), cmap=label_camp)  # 第二张图像，设置透明度
    plt.title(f"Slice at Z = {z_index}")
    plt.axis("off")
    plt.show()


def plot_image(image, label, z_index):
    # 提取二维切片
    img_array, label_array = extract_slice(image, label, z_index)

    # 绘制二维切片
    plot_slice(img_array, label_array, z_index)


def main():
    # 用于存储所有 image 和 label 数据的列表
    image_list = []
    label_list = []

    for i in range(1, 1063):
        filename = f"{i:05d}.mha"

        # 加载图像
        image_res = load_image(filename)
        if image_res is not None:
            image_array, label_array, img_shape = image_res

            # 存入列表
            image_list.append(image_array)
            label_list.append(label_array)

    # 转换为 NumPy 数组（最终 shape: (N, Z, Y, X)）
    image_data  = np.stack(image_list, axis=0)
    label_data = np.stack(label_list, axis=0)

    print(f"Final Image Shape: {image_data.shape}")
    print(f"Final Label Shape: {label_data.shape}")

    plot_image(image_list[0], label_list[0], 80)


if __name__ == "__main__":
    main()
